"""
This script is for the backtesting of signals.
"""
import pandas as pd
import matplotlib.pyplot as plt
from typing import List, Union
from joblib import Parallel, delayed

from .factor import *
from .factor_indicators import FactorIndicator
from .factor_utils import get_ret
from lwpackage.lwdata import get_stock_price, get_futures_continuous_contract_price
from lwpackage.lwutils.logging import log
from lwpackage.lwerror.errors import NotBackTestingError


class BaseBackTester:
    """
    A class for single signal backtesting.
    """

    def __init__(self,
                 data: Union[pd.DataFrame, None] = None,
                 instrument_type: str = 'stock',
                 fc_number: int = None,
                 fc_name: str = None,
                 start_time: str = None,
                 end_time: str = None,
                 ret_freq: int = 30,
                 transaction_period: Union[int, None] = None,
                 fc_freq: str = '1d',
                 risk_free_rate: bool = True,
                 cost: Union[int, None] = 0.00025,
                 quantile: float = 0.2,
                 method: str = 'longshort',
                 ir: bool = True,
                 longshort_instrument_number: int = 5,
                 yearly_performance: bool = True,
                 preprocess: bool = True,
                 backtest: bool = False):
        """
        Some params need to be initialized for backtesting. We assume that the signal is generated at
        the open time of every bar, and we immediately trade at this bar, and will complete the transaction
        at the close time of the next `transaction_period` bar with average transaction price
        rolling_mean((open + high + low + close) / 4)

        :param data: data should be a dataframe with factors of all instrument in every bar.
            If data is provided, the following process of using fc_number to get factor value will be ignored.
        :param fc_number: the number of factor
        :param start_time: backtesting start time,
            default is the earliest price data that can be found on database
        :param end_time: backtesting end time,
            default is the latest price data that can be found on database
        :param ret_freq: we calculate return as the return between current bar and the next `ret_freq` bars
        :param transaction_period: the period for one transaction to complete
        :param fc_freq: the frequency of factor, 1m, 5m or 1d
        :param risk_free_rate: whether to consider risk-free rate when calculating excess return
            as well as sharpe ratio
        :param cost: the cost of trade, if None, will not calculate any net values
        :param quantile: the quantile for constructing long-short portfolio. default is to 0.2, meaning we
            sort underlying from max to min according to the factor value on every time cross-section.
            And we use top 20% underlying to construct long position; bottom 20% underlying for
            short position. Also, we construct a capital-equalized portfolio.
        :param method: the method for calculating sharpe, can be 'longshort' or 'long_only' or 'zscore',
            default is 'longshort'
        :param ir: if ir, will calculate icir
        :param longshort_instrument_number: number of long and short lines for the instrument with large value and small
            value
        :param yearly_performance: if True, performance will include yearly performance,
            otherwise only total performance.
        :param preprocess: if True, will preprocess the stock data, which takes longer for initializing
            this class. If data is provided, preprocess will have no effect.
        """
        assert fc_freq in ['1m', '5m', '1d'], f'Only support 1m, 5m or 1d fc_freq, but got {fc_freq} instead.'

        self.data = data
        self.instrument_type = instrument_type
        self.fc_number = fc_number
        self.fc_name = fc_name
        self.start_time = start_time
        self.end_time = end_time
        self.ret_freq = ret_freq
        self.transaction_period = transaction_period
        self.fc_freq = fc_freq
        self.rfr = risk_free_rate
        self.cost = cost
        self.quantile = quantile
        self.method = method
        self.ir = ir
        self.longshort_instrument_number = longshort_instrument_number
        self.yearly_performance = yearly_performance

        if isinstance(self.data, pd.DataFrame):
            assert self.data.index.names[0] == 'time', \
                f'Invalid data format, need `time` to be the first of multi-index.'
            assert self.data.index.names[1] == 'instrument_id', \
                f'Invalid data format, need `instrument_id` to be the second of multi-index.'

            assert not self.fc_number, f'Data has been provided, should not specify factor number.'
            assert self.fc_name, f'Please specify the name of this factor generating method.'
            self.fc_number = 'fc_from_model'
            self.is_preprocessed = True
        else:
            if self.instrument_type == 'stock':
                self.data = get_stock_price(frequency=self.fc_freq,
                                            start_time=self.start_time,
                                            end_time=self.end_time,
                                            from_database=True)
            elif self.instrument_type == 'futures_continuous_contract':
                self.data = get_futures_continuous_contract_price(start_date=self.start_time,
                                                                  end_date=self.end_time,
                                                                  from_database=True)
            else:
                raise ValueError(f'Does not support instrument type {self.instrument_type}.')
            self.is_preprocessed = False

        self.preprocess = preprocess
        self.is_backtested = False

        if self.preprocess:
            self._preprocess_data()

        # the performance summary
        self.performance = None
        # all factor indicators, derived from FactorIndicator class.
        self.ai = None

        if backtest:
            self.backtest()

    def _generate_fc_value(self):
        """
        Generate factor value.
        self.data is a panel dataframe containing all instrument info of all time.
        """

        assert self.data.index.names[0] == 'time'
        assert self.data.index.names[1] == 'instrument_id'

        mapper = (
            self.data
            .groupby('instrument_id', as_index=False, sort=False)
            .apply(eval(f'factormaker{self.fc_number}'))
            .reset_index().set_index(['time', 'instrument_id']).iloc[:, -1]
        )

        if mapper.name in self.data.columns:
            self.data.drop(columns=mapper.name, inplace=True)
        self.data[f'{mapper.name}'] = self.data.index.map(mapper)
        # Attention: if isinstance(self.data, pd.DataFrame), here is never reached.
        self.fc_name = mapper.name

    def _standardize_cross_sectional_fc(self):
        """
        Standardize the factor value on every time cross-section.
        """
        assert 'time' == self.data.index.names[0]
        assert 'signal' in self.data.columns.get_level_values(0)

        mapper = self.data['signal'].groupby('time').apply(lambda x: (x - x.mean()) / x.std()).droplevel(0)
        self.data.loc[:, 'signal'] = self.data.index.map(mapper)

    def _preprocess_data(self):
        """
        Preprocess stock data for backtesting.
        """
        if self.is_preprocessed:
            return

        # reindex
        for col in ['time', 'instrument_id']:
            assert col in self.data.columns, f'self.data does not contain column {col}.'
        self.data.set_index(['time', 'instrument_id'], inplace=True)

        # fc value
        self._generate_fc_value()

        # get return as label
        self.data = get_ret(self.data, self.fc_freq, self.transaction_period, self.ret_freq, self.rfr)

        # missing value
        self.data.dropna(how='any', axis=0, inplace=True)

        # rename columns
        assert self.fc_name is not None
        self.data = self.data[[self.fc_name, 'ret']].copy()
        self.data.columns = ['signal', 'ret']

        # standardize factor values to zscore
        self._standardize_cross_sectional_fc()

        self.is_preprocessed = True

    def plot_nav(self, net: bool = True):
        """
        Plot the cumulative return(i.e. nav) graph for factor.
        e.g. bt.plot_nav()

        :param net: if net and considering cost, then draw the graph for both gross and net nav.
        :return: None
        """

        if not self.is_backtested:
            raise NotBackTestingError('Need to backtest first before plotting nav.')

        if self.cost and net:
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 16))
            ax1.plot(self.ai.ret_daily['gross_ret'].index, self.ai.ret_daily['gross_ret'].cumsum())
            ax1.set_title(f'Cumulative Gross Return of factor {self.fc_number}')
            ax1.set_xlabel('time')
            ax2.plot(self.ai.ret_daily['net_ret'].index, self.ai.ret_daily['gross_ret'].cumsum())
            ax2.set_title(f'Cumulative Net Return of factor {self.fc_number}')
            ax2.set_xlabel('time')
        else:
            fig, ax = plt.subplots(figsize=(10, 10))
            ax.plot(self.ai.ret_daily['gross_ret'].index, self.ai.ret_daily['gross_ret'].cumsum())
            ax.set_title(f'Cumulative Return of factor {self.fc_number}')
            ax.set_xlabel('time')

    def plot_longshort(self):
        """
        Plot the ret for n-largest and n-smallest signal values, with n equals to self.longshort_instrument_number.
        e.g. the factor value from ret_short5, ret_short_4, ..., ret_short_1, ret_long_1, ret_long_2, ..., ret_long_5
            is becoming larger.
        e.g. bt.plot_longshort()
        """
        if not self.is_backtested:
            raise NotBackTestingError('Need to backtest first before plotting longshort.')

        fig, ax = plt.subplots(figsize=(12, 16))
        for col in self.ai.longshort_ret.columns:
            ax.plot(self.ai.longshort_ret.index, self.ai.longshort_ret[col].cumsum(), label=col)
        ax.legend()
        ax.set_title(f'Longshort ret of factor {self.fc_number}')
        ax.set_xlabel('time')

    def backtest(self):
        """
        use this method for backtesting.
        e.g. bt = BaseBackTester(1)
             bt.backtest()
        Then the performance will be stored on bt.performance. All factors will be stored on bt.ai
        """
        if not self.is_preprocessed:
            self._preprocess_data()

        fi = FactorIndicator(data=self.data,
                             cost=self.cost,
                             fc_freq=self.fc_freq,
                             performance=True,
                             quantile=self.quantile,
                             method=self.method,
                             ir=self.ir,
                             longshort_instrument_number=self.longshort_instrument_number)
        if self.yearly_performance:
            self.performance = fi.performance.copy()
        else:
            self.performance = fi.performance_all.copy()

        temp_cols = self.performance.columns.to_list()
        assert temp_cols[0] == 'year', f'{temp_cols}'
        self.performance.loc[:, 'fc_number'] = self.fc_number
        self.performance.loc[:, 'fc_name'] = self.fc_name
        self.performance.loc[:, 'fc_freq'] = self.fc_freq
        indicator_cols = temp_cols[1:]
        self.performance = self.performance[['year', 'fc_number', 'fc_name', 'fc_freq'] + indicator_cols].copy()
        for col in indicator_cols:
            self.performance[col] = self.performance[col].apply(lambda x: '{:.6f}'.format(x))
        # ai means all indicators
        self.ai = fi
        self.is_backtested = True
        log.info(f'Successfully generate backtesting result for factor {self.fc_number} '
                 f'on data with frequency {self.fc_freq}')


class BackTester:
    """
    A class supports multi-signal backtesting using Parallel.
    """

    def __init__(self,
                 data: Union[pd.DataFrame, None] = None,
                 instrument_type: str = 'stock',
                 fc_name: str = None,
                 fc_number_list: Union[int, List] = None,
                 start_time: str = None,
                 end_time: str = None,
                 ret_freq: int = 30,
                 transaction_period: Union[int, None] = None,
                 fc_freq: str = '1d',
                 risk_free_rate: bool = True,
                 cost: Union[int, None] = 0.00025,
                 quantile: float = 0.2,
                 method: str = 'longshort',
                 ir: bool = True,
                 longshort_instrument_number: int = 5,
                 yearly_performance: bool = True,
                 n_jobs: int = -1):
        """
        Some params need to be initialized for backtesting.

        :param data: data should be a dataframe with factors of all instrument in every bar. The factor values
            need to be preprocessed. If given data, there will only be one factor.
        :param fc_name: If given data, should also specify fc_name, otherwise fc_name will be the name of factorMaker
            return's factor column name.
        :param fc_number_list: the list of factor numbers
        :param start_time: backtesting start time,
            default is the earliest price data that can be found on database
        :param end_time: backtesting end time,
            default is the latest price data that can be found on database
        :param ret_freq: we use ret data which is ret_freq bars after the factor to calculate ic, icir, etc.
        :param transaction_period: the period for one transaction to complete
        :param fc_freq: the frequency of factor, 1m, 5m or 1d.
        :param risk_free_rate: whether to consider risk-free rate when calculating excess return
            as well as sharpe ratio
        :param cost: the cost of trade, if None, will not calculate any net values.
        :param quantile: the quantile for constructing long-short portfolio. default is to 0.2, meaning we
            sort underlying from max to min according to the factor value on every time cross-section.
            And we use top 20% underlying to construct long position; bottom 20% underlying for
            short position. Also, we construct a capital-equalized portfolio.
        :param method: the method for calculating sharpe, can be 'longshort' or 'long_only' or 'zscore',
            default is 'longshort'
        :param ir: if ir, will calculate icir.
        :param longshort_instrument_number: number of long and short lines for the instrument with large value and small
            value
        :param yearly_performance: if True, performance will include yearly performance,
            otherwise only total performance
        :param n_jobs: Parallel's n_job param. Default is to use all available CPUs (-1)
        """
        self.data = data
        self.instrument_type = instrument_type
        self.fc_name = fc_name
        self.fc_number_list = fc_number_list
        self.start_time = start_time
        self.end_time = end_time
        self.ret_freq = ret_freq
        self.transaction_period = transaction_period
        self.fc_freq = fc_freq
        self.rfr = risk_free_rate
        self.cost = cost
        self.quantile = quantile
        self.method = method
        self.ir = ir
        self.longshort_instrument_number = longshort_instrument_number
        self.yearly_performance = yearly_performance
        self.n_jobs = n_jobs

        self.performance = None
        self.ai = None
        self.bbt = None

        if isinstance(self.data, pd.DataFrame):
            assert not self.fc_number_list, f'When specifying data, should not specify fc_number_list.'
        # we allow fc_number to be inconsistent.
        elif isinstance(self.fc_number_list, int):
            self.fc_number_list = [self.fc_number_list]
        else:
            assert isinstance(self.fc_number_list, list)
            self.fc_number_list = self.fc_number_list

        self.is_backtested = False

    def plot_nav(self, fc_number: int, net: bool = True):
        """
        Plot the cumulative return(i.e. nav) graph for given factor.
        e.g. bt.plot_nav(1)

        :param fc_number: factor number.
        :param net: if net and considering cost, then draw the graph for both gross and net nav.
        :return: None
        """
        if not self.bbt[fc_number].is_backtested:
            raise NotBackTestingError('Need to backtest first before plotting nav.')
        self.bbt[fc_number].plot_nav(net=net)

    def plot_longshort(self, fc_number: int):
        """
        Plot the ret for n-largest and n-smallest signal values for given factor, with n equals to
        self.longshort_instrument_number.
        e.g. bt.plot_longshort(1)

        :param fc_number: factor number.
        :return: None
        """
        if not self.bbt[fc_number].is_backtested:
            raise NotBackTestingError('Need to backtest first before plotting longshort.')
        self.bbt[fc_number].plot_longshort()

    def backtest(self):
        """
        Use this method for backtesting.
        e.g. bt = BackTester([1, 2])
             bt.backtest()
        Then the performance will be stored on bt.performance.
        bt.ai is a dict which stores the performance of `all factors`. The key is the factor number.
        For example. bt.ai[1] stores all indicators of factor1. Using bt.ai[1].turnover to get turnover.
        """

        # bt is the list of BaseBackTesters for factor number in self.fc_number_list
        f = lambda x: BaseBackTester(self.data, self.instrument_type, x, self.fc_name, self.start_time,
                                     self.end_time, self.ret_freq, self.transaction_period, self.fc_freq,
                                     self.rfr, self.cost, self.quantile, self.method, self.ir,
                                     self.longshort_instrument_number, self.yearly_performance,
                                     preprocess=True, backtest=True)
        if isinstance(self.data, pd.DataFrame):
            self.bbt = f(None)
            performance_list = [self.bbt.performance]
            all_indicators_dc = {self.bbt.fc_number: self.bbt.ai}
        else:
            bt = Parallel(n_jobs=self.n_jobs)(
                delayed(f)(fc_number) for fc_number in self.fc_number_list)
            self.bbt = {k: v for k, v in zip(self.fc_number_list, bt)}

            performance_list = [self.bbt[fc_number].performance for fc_number in self.fc_number_list]

            all_indicators_dc = {fc_number: self.bbt[fc_number].ai for fc_number in self.fc_number_list}

        self.performance = pd.concat(performance_list)
        # e.g. self.ai[1] is self.ai for factor 1 of BaseBackTester
        self.ai = all_indicators_dc

        self.is_backtested = True
        log.info(f'Successfully generate backtesting result for all factors '
                 f'on data with frequency {self.fc_freq}')
